from typing import List, Dict
import time
from app.models.llm import TinyLlamaModel, TinyLlamaLLM
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain

class ChatService:
    """聊天服务类"""
    
    def __init__(self):
        self.model = TinyLlamaModel()
        self.llm = TinyLlamaLLM(self.model)
        self.conversations: Dict[str, ConversationChain] = {}
        
    def get_or_create_conversation(self, session_id: str) -> ConversationChain:
        """获取或创建会话"""
        if session_id not in self.conversations:
            memory = ConversationBufferMemory()
            self.conversations[session_id] = ConversationChain(
                llm=self.llm,
                memory=memory,
                verbose=True
            )
        return self.conversations[session_id]
        
    async def chat(self, session_id: str, message: str) -> str:
        """处理聊天消息
        
        Args:
            session_id: 会话ID
            message: 用户消息
            
        Returns:
            模型回复
        """
        conversation = self.get_or_create_conversation(session_id)
        try:
            response = await conversation.apredict(input=message)
            return response
        except Exception as e:
            raise RuntimeError(f"聊天失败: {str(e)}") 